ARG BASE_IMAGE=pytorch/pytorch:2.5.1-cuda11.8-cudnn9-runtime
FROM ${BASE_IMAGE}

ENV ALGO_NAME=GLEE
ENV APP_ROOT=/opt/${ALGO_NAME}
ENV CUDA_HOME=/usr/local/cuda \
    TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6+PTX" \
    SETUPTOOLS_USE_DISTUTILS=stdlib
ENV DEBIAN_FRONTEND=noninteractive

# 预先配置时区
RUN ln -fs /usr/share/zoneinfo/Asia/Shanghai /etc/localtime && \
    echo "Asia/Shanghai" > /etc/timezone
RUN apt-get update && apt-get install -y \
    python3-opencv ca-certificates python3-dev git wget ninja-build build-essential gcc sudo

# 升级 pip 和 setuptools
RUN pip install --upgrade pip setuptools --index-url https://pypi.tuna.tsinghua.edu.cn/simple

# 创建应用目录
RUN mkdir -p ${APP_ROOT}

# 复制 GLEE 代码到容器中
COPY ./models/${ALGO_NAME} ${APP_ROOT}/agents/models/${ALGO_NAME}

# 复制 mq 代码到容器中
COPY ./mq ${APP_ROOT}/agents/mq

# 设置工作目录
WORKDIR ${APP_ROOT}/agents/models/${ALGO_NAME}

# 下载github GroundingDINO仓库
RUN git clone https://github.com/FoundationVision/GLEE.git
RUN cd GLEE/ && python -m pip install .

# 安装 Python 依赖，使用清华源
RUN pip install -r requirements.txt --no-deps --no-build-isolation -v --index-url https://pypi.tuna.tsinghua.edu.cn/simple
# RUN pip install git+https://github.com/wjf5203/cocoapi.git#"egg=pycocotools&subdirectory=PythonAPI"

# 下载模型文件到对应的目录
RUN mkdir -p ${APP_ROOT}/checkpoints
ADD https://huggingface.co/spaces/Junfeng5/GLEE_demo/resolve/main/MODEL_ZOO/GLEE_Lite_joint.pth ${APP_ROOT}/checkpoints/GLEE_Lite_joint.pth
ADD https://huggingface.co/spaces/Junfeng5/GLEE_demo/resolve/main/MODEL_ZOO/GLEE_Plus_joint.pth ${APP_ROOT}/checkpoints/GLEE_Plus_joint.pth
ADD https://huggingface.co/spaces/Junfeng5/GLEE_demo/resolve/main/MODEL_ZOO/GLEE_Pro_joint.pth ${APP_ROOT}/checkpoints/GLEE_Pro_joint.pth
ADD https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window12_384_22k.pth ${APP_ROOT}/checkpoints/swin_large_patch4_window12_384_22k.pth

# Download pretrained Language Model (CLIP text encoder from transformers)
ADD https://huggingface.co/spaces/Junfeng5/GLEE_demo/resolve/main/GLEE/clip_vit_base_patch32/pytorch_model.bin ${APP_ROOT}/projects/GLEE/clip_vit_base_patch32/pytorch_model.bin

# 修改detectron2的问题（升级pillow库后，原来文件不兼容，更新下代码）
COPY ./models/${ALGO_NAME}/patch/transform.py /opt/conda/lib/python3.11/site-packages/detectron2/data/transforms/transform.py

# 设置启动命令
CMD ["python", "run.py"]